import cv2
import numpy as np


# sigmoid
def sigmoid(x):
    s = 1 / (1 + np.exp(-x))
    return s


# 画出图像中信息
def draw_image(img, results, class_names):
    if results[0] is not None:
        label = np.array(results[0][:, 6], dtype='int32')
        conf = results[0][:, 4] * results[0][:, 5]
        boxes = results[0][:, :4]
    for i, c in list(enumerate(label)):
        predicted_class = class_names[int(c)]
        box = boxes[i]
        score = conf[i]

        label = '{} {:.2f}'.format(predicted_class, score)
        top, left, bottom, right = box
        top = max(0, np.floor(top).astype('int32'))
        left = max(0, np.floor(left).astype('int32'))
        bottom = min(img.shape[0], np.floor(bottom).astype('int32'))
        right = min(img.shape[1], np.floor(right).astype('int32'))
        cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 2)
        cv2.putText(img, str(label), (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    return img


# 网络预测信息解码过程
def decode(inputs, num_classes, input_shape=[416, 416], anchor_nums=3, anchor_mask=[[3, 4, 5], [0, 1, 2], [6, 7, 8]],
           anchors=np.array(
               [[10, 13.], [16, 30.], [33., 23.], [30, 61.], [62., 45.], [59., 119.], [116., 90.], [156., 198.],
                [373., 326.]])):
    """
        :param outputs: outputs.size = (batch_size, anchor_nums*(num_classes + 5),input_size.h,input_size.w)
        :param anchors:anchors = [[ 10.  13.][ 16.  30.][ 33.  23.][ 30.  61.][ 62.  45.][ 59. 119.][116.  90.][156. 198.][373. 326.]]
        :param anchor_nums: anchor_nums = 3
        :param num_classes: num_classes in the VOC dates equal 20 and in the COCO dates equal 80
        :return:(N,num_classes + 5) N is the numbers of prediction boxs, and 5 is the x,y,w,h and num_classes, the numbers of num_classes is 20 in the VOC dates
        """
    result = []
    for i, input in enumerate(inputs):
        batch_size = input.shape[0]
        input_height = input.shape[2]
        input_width = input.shape[3]
        # (batch_size,(num-classes + 5)* anchor_nums,N,N) -> (batch_size,anchor_nums,N,N,num_classes + 5)
        prediction = input.reshape(batch_size, anchor_nums, num_classes + 5, input_height, input_width) \
            .transpose(0, 1, 3, 4, 2)

        # tx and ty
        tx = sigmoid(prediction[..., 0])
        ty = sigmoid(prediction[..., 1])
        # w and h
        w = prediction[..., 2]
        h = prediction[..., 3]
        # confidence and classes
        conf = sigmoid(prediction[..., 4])
        pre_cls = sigmoid(prediction[..., 5:])

        # create grid x and grid y
        # batch_size, anchor_nums, N,N
        grid_x = np.linspace(0, input_width - 1, input_width).reshape(1, input_width).repeat(input_height, 0) \
            .reshape(1, input_width, input_height).repeat(anchor_nums, 0).reshape(batch_size, anchor_nums, input_height,
                                                                                  input_width)
        grid_y = np.linspace(0, input_height - 1, input_height).reshape(1, input_height).repeat(input_width, 0).T \
            .reshape(1, input_width, input_height).repeat(anchor_nums, 0).reshape(batch_size, anchor_nums, input_height,
                                                                                  input_width)

        # 每个网格点在原图上占多少个网格
        stride_h = input_shape[0] / input_height
        stride_w = input_shape[1] / input_width

        # 每个特征层的边长的大小
        scaled = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in
                  anchors[anchor_mask[i]]]

        #   按照网格格式生成先验框的宽高
        #   batch_size, anchor_nums, N,N
        anchors_w = np.expand_dims(np.expand_dims(np.array(scaled)[:, 0], axis=1)
                                   .repeat(input_width, axis=1), axis=1).repeat(input_height, axis=1) \
            .reshape(batch_size, anchor_nums, input_height, input_width)
        anchors_h = np.expand_dims(np.expand_dims(np.array(scaled)[:, 1], axis=1)
                                   .repeat(input_width, axis=1), axis=1).repeat(input_height, axis=1) \
            .reshape(batch_size, anchor_nums, input_height, input_width)

        prediction[..., 0] = tx + grid_x
        prediction[..., 1] = ty + grid_y
        prediction[..., 2] = np.exp(w) * anchors_w
        prediction[..., 3] = np.exp(h) * anchors_h

        # 归一化
        _scale = np.array([input_width, input_height, input_width, input_height], dtype=np.float32)
        prediction[..., :4] = prediction[..., :4] / _scale

        prediction[..., 4] = conf
        prediction[..., 5:] = pre_cls

        prediction = prediction.reshape(-1, num_classes + 5)
        result.append(prediction)
    result = np.concatenate(result).reshape(batch_size, -1, num_classes + 5)

    return result


# 将图像中的回归预测结果正确解码后添加
def yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image=False):
    # -----------------------------------------------------------------#
    #   把y轴放前面是因为方便预测框和图像的宽高进行相乘
    # -----------------------------------------------------------------#
    box_yx = box_xy[..., ::-1]
    box_hw = box_wh[..., ::-1]
    input_shape = np.array(input_shape)
    image_shape = np.array(image_shape)

    if letterbox_image:
        # -----------------------------------------------------------------#
        #   这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
        #   new_shape指的是宽高缩放情况
        # -----------------------------------------------------------------#
        new_shape = np.round(image_shape * np.min(input_shape / image_shape))
        offset = (input_shape - new_shape) / 2. / input_shape
        scale = input_shape / new_shape

        box_yx = (box_yx - offset) * scale
        box_hw *= scale
    box_mins = box_yx - (box_hw / 2.)
    box_maxes = box_yx + (box_hw / 2.)
    boxes = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]],
                           axis=-1)
    boxes *= np.concatenate([image_shape, image_shape], axis=-1)
    return boxes


# 非极大值抑制
def non_max_suppression(prediction, num_classes, image_shape, input_shape=[416, 416], conf_thres=0.5, nms_thres=0.5):
    # ----------------------------------------------------------#
    #   将预测结果的格式转换成左上角右下角的格式。
    #   prediction  [num_anchors, num_classes + 5]
    # ----------------------------------------------------------#
    box_corner = np.zeros_like(prediction)
    box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
    box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
    prediction[:, :, :4] = box_corner[:, :, :4]

    prediction[:, :4] = box_corner[:, :4]

    output = [None for _ in range(len(prediction))]
    for i, image_pred in enumerate(prediction):
        # ----------------------------------------------------------#
        #   对种类预测部分取max。
        #   class_conf  [num_anchors, 1]    种类置信度
        #   class_pred  [num_anchors, 1]    种类
        # ----------------------------------------------------------#
        class_conf = np.max(image_pred[:, 5:5 + num_classes], 1, keepdims=True)
        class_pred = np.argmax(image_pred[:, 5:5 + num_classes], 1, keepdims=True)

        # ----------------------------------------------------------#
        #   利用置信度进行第一轮筛选
        # ----------------------------------------------------------#
        conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()

        # ----------------------------------------------------------#
        #   根据置信度进行预测结果的筛选
        # ----------------------------------------------------------#
        image_pred = image_pred[conf_mask]
        class_conf = class_conf[conf_mask]
        class_pred = class_pred[conf_mask]

        if not image_pred.shape[0]:
            continue
        # -------------------------------------------------------------------------#
        #   detections  [num_anchors, 7]
        #   7的内容为：x1, y1, x2, y2, obj_conf, class_conf, class_pred
        # -------------------------------------------------------------------------#
        detections = np.concatenate((image_pred[:, :5], class_conf, class_pred), 1)

        # ------------------------------------------#
        #   获得预测结果中包含的所有种类
        # ------------------------------------------#
        unique_labels = np.unique(detections[:, -1])

        for c in unique_labels:
            # ------------------------------------------#
            #   获得某一类得分筛选后全部的预测结果
            # ------------------------------------------#
            detections_class = detections[detections[:, -1] == c]
            keep = cv2.dnn.NMSBoxes(detections_class[:, :4], detections_class[:, 4] * detections_class[:, 5],
                                    conf_thres, nms_thres)
            max_detections = detections_class[keep]
            output[i] = max_detections if output[i] is None else np.concatenate((output[i], max_detections))
        if output[i] is not None:
            box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4]) / 2, output[i][:, 2:4] - output[i][:, 0:2]
            output[i][:, :4] = yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape)

        return output


# 定义标签名自动读取
def read_label_names(path):
    # 读取每一行数据并获取类别信息
    with open(path, 'r') as f:
        class_names = [line.strip() for line in f.readlines()]
    return class_names


# 不失真的resize
def letterbox_image(img, size):
    # 对图片进行resize,使得图像不失真，在空缺的地方进行padding
    iw, ih = img.shape[:2]
    w, h = size[:2]
    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)
    new_image = np.zeros(size, dtype='uint8')
    img = cv2.resize(img, (nh, nw))
    if (w - nw) % 2 != 0:
        new_image[((w - nw) // 2):w - ((w - nw) // 2) - 1, ((h - nh) // 2):h - ((h - nh) // 2), :] = img
    elif (h - nh) % 2 != 0:
        new_image[((w - nw) // 2):w - ((w - nw) // 2), ((h - nh) // 2):h - ((h - nh) // 2) - 1, :] = img
    elif (w - nw) % 2 != 0 and (h - nh) % 2 != 0:
        new_image[((w - nw) // 2) - 1:w - ((w - nw) // 2), ((h - nh) // 2):h - ((h - nh) // 2) - 1, :] = img
    else:
        new_image[((w - nw) // 2):w - ((w - nw) // 2), ((h - nh) // 2):h - ((h - nh) // 2), :] = img
    return new_image


if __name__ == '__main__':
    img = cv2.imread('../img/img.png')
    img = letterbox_image(img, [416, 416, 3])
    cv2.imshow('1', img)
    cv2.waitKey(0)
    print(img.shape)
    print()
